// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <melon/base64.h>
#include <pollux/testing/gtest_utils.h>
#include <pollux/functions/lib/tdigest.h>
#include <pollux/testing/functions/prestosql/function_base_test.h>
#include <pollux/functions/prestosql/types/tdigest_registration.h>
#include <pollux/functions/prestosql/types/tdigest_type.h>

using namespace kumo::pollux;
using namespace kumo::pollux::exec;
using namespace kumo::pollux::functions::test;

class TDigestFunctionsTest : public FunctionBaseTest {
 protected:
  void SetUp() override {
    FunctionBaseTest::SetUp();
    registerTDigestType();
  }
  std::string decodeBase64(std::string_view input) {
    std::string decoded(melon::base64DecodedSize(input), '\0');
    melon::base64Decode(input, decoded.data());
    return decoded;
  }
  const TypePtr TDIGEST_DOUBLE = TDIGEST(DOUBLE());
  const TypePtr ARRAY_TDIGEST_DOUBLE = ARRAY(TDIGEST(DOUBLE()));
  // Default compression and weight
  // Digest 1 has one value of 0.1
  const std::string digest1String = decodeBase64(
      "AQCamZmZmZm5P5qZmZmZmbk/mpmZmZmZuT8AAAAAAABZQAAAAAAAAPA/AQAAAAAAAAAAAPA/mpmZmZmZuT8=");
  // Digest 2 has one value of 0.2
  const std::string digest2String = decodeBase64(
      "AQCamZmZmZnJP5qZmZmZmck/mpmZmZmZyT8AAAAAAABZQAAAAAAAAPA/AQAAAAAAAAAAAPA/mpmZmZmZyT8=");
  // Digest 12 has merged Digest 1 and Digest 2
  const std::string digest12String = decodeBase64(
      "AQCamZmZmZm5P5qZmZmZmck/NDMzMzMz0z8AAAAAAABZQAAAAAAAAABAAgAAAAAAAAAAAPA/AAAAAAAA8D+amZmZmZm5P5qZmZmZmck/");
};

TEST_F(TDigestFunctionsTest, valueAtQuantile) {
  const auto valueAtQuantile = [&](const std::optional<std::string>& input,
                                   const std::optional<double>& quantile) {
    return evaluateOnce<double>(
        "value_at_quantile(c0, c1)", TDIGEST_DOUBLE, input, quantile);
  };
  const std::string input = decodeBase64(
      "AQAAAAAAAADwPwAAAAAAABRAAAAAAAAALkAAAAAAAABZQAAAAAAAABRABQAAAAAAAAAAAPA/AAAAAAAA8D8AAAAAAADwPwAAAAAAAPA/AAAAAAAA8D8AAAAAAADwPwAAAAAAAABAAAAAAAAACEAAAAAAAAAQQAAAAAAAABRA");
  ASSERT_EQ(1.0, valueAtQuantile(input, 0.1));
  ASSERT_EQ(3.0, valueAtQuantile(input, 0.5));
  ASSERT_EQ(5.0, valueAtQuantile(input, 0.9));
  ASSERT_EQ(5.0, valueAtQuantile(input, 0.99));
};

TEST_F(TDigestFunctionsTest, valuesAtQuantiles) {
  const std::string input = decodeBase64(
      "AQAAAAAAAADwPwAAAAAAABRAAAAAAAAALkAAAAAAAABZQAAAAAAAABRABQAAAAAAAAAAAPA/AAAAAAAA8D8AAAAAAADwPwAAAAAAAPA/AAAAAAAA8D8AAAAAAADwPwAAAAAAAABAAAAAAAAACEAAAAAAAAAQQAAAAAAAABRA");
  auto arg0 = make_flat_vector<std::string>({input}, TDIGEST_DOUBLE);
  auto arg1 = make_nullable_array_vector<double>({{0.1, 0.5, 0.9, 0.99}});
  auto expected = make_nullable_array_vector<double>({{1.0, 3.0, 5.0, 5.0}});
  auto result =
      evaluate("values_at_quantiles(c0, c1)", make_row_vector({arg0, arg1}));
  test::assertEqualVectors(expected, result);
}

TEST_F(TDigestFunctionsTest, testMergeTDigestNullInput) {
  auto arg0 = make_nullable_array_vector<std::string>(
      {std::nullopt}, ARRAY_TDIGEST_DOUBLE);
  auto expected =
      make_nullable_flat_vector<std::string>({std::nullopt}, TDIGEST_DOUBLE);
  auto result = evaluate("merge_tdigest(c0)", make_row_vector({arg0}));
  test::assertEqualVectors(expected, result);
}

TEST_F(TDigestFunctionsTest, testMergeTDigestEmptyArray) {
  auto arg0 =
      vectorMaker_.array_vector_nullable<std::string>({{}}, ARRAY_TDIGEST_DOUBLE);
  auto expected =
      make_nullable_flat_vector<std::string>({std::nullopt}, TDIGEST_DOUBLE);
  auto result = evaluate("merge_tdigest(c0)", make_row_vector({arg0}));
  test::assertEqualVectors(expected, result);
}

TEST_F(TDigestFunctionsTest, testMergeTDigestEmptyArrayOfNull) {
  auto arg0 = vectorMaker_.array_vector_nullable<std::string>(
      {{std::nullopt}}, ARRAY_TDIGEST_DOUBLE);
  auto expected =
      make_nullable_flat_vector<std::string>({std::nullopt}, TDIGEST_DOUBLE);
  auto result = evaluate("merge_tdigest(c0)", make_row_vector({arg0}));
  test::assertEqualVectors(expected, result);
}

TEST_F(TDigestFunctionsTest, testMergeTDigestEmptyArrayOfNulls) {
  auto arg0 = vectorMaker_.array_vector_nullable<std::string>(
      {{{std::nullopt, std::nullopt, std::nullopt}}}, ARRAY_TDIGEST_DOUBLE);
  auto expected =
      make_nullable_flat_vector<std::string>({std::nullopt}, TDIGEST_DOUBLE);
  auto result = evaluate("merge_tdigest(c0)", make_row_vector({arg0}));
  test::assertEqualVectors(expected, result);
}

TEST_F(TDigestFunctionsTest, testMergeTDigests) {
  auto arg0 = make_nullable_array_vector<std::string>(
      {{digest1String, digest2String}}, ARRAY_TDIGEST_DOUBLE);
  auto expected =
      make_nullable_flat_vector<std::string>({digest12String}, TDIGEST_DOUBLE);
  auto result = evaluate("merge_tdigest(c0)", make_row_vector({arg0}));
  test::assertEqualVectors(expected, result);
}

TEST_F(TDigestFunctionsTest, testMergeTDigestOneNull) {
  auto arg0 = make_nullable_array_vector<std::string>(
      {{digest1String, std::nullopt}}, ARRAY_TDIGEST_DOUBLE);
  auto expected =
      make_nullable_flat_vector<std::string>({digest1String}, TDIGEST_DOUBLE);
  auto result = evaluate("merge_tdigest(c0)", make_row_vector({arg0}));
  test::assertEqualVectors(expected, result);
}

TEST_F(TDigestFunctionsTest, testMergeTDigestOneNullFirst) {
  auto arg0 = make_nullable_array_vector<std::string>(
      {{std::nullopt, digest1String, digest2String}}, ARRAY_TDIGEST_DOUBLE);
  auto expected =
      make_nullable_flat_vector<std::string>({digest12String}, TDIGEST_DOUBLE);
  auto result = evaluate("merge_tdigest(c0)", make_row_vector({arg0}));
  test::assertEqualVectors(expected, result);
}

TEST_F(TDigestFunctionsTest, testMergeTDigestOneNullMiddle) {
  auto arg0 = make_nullable_array_vector<std::string>(
      {{digest1String, std::nullopt, digest2String}}, ARRAY_TDIGEST_DOUBLE);
  auto expected =
      make_nullable_flat_vector<std::string>({digest12String}, TDIGEST_DOUBLE);
  auto result = evaluate("merge_tdigest(c0)", make_row_vector({arg0}));
  test::assertEqualVectors(expected, result);
}
